"""交通监控行为分析管理器模块"""

import cv2
import numpy as np
from typing import Dict, List, Tuple, Optional
import time
from config import get_config
from .traffic_analyzer import TrafficAnalyzer
from .violation_detector import ViolationDetector

class AnalysisManager:
    """交通监控行为分析管理器，协调各种交通场景分析功能"""
    
    def __init__(self):
        """初始化交通监控行为分析管理器"""
        # 初始化分析器和检测器
        self.traffic_analyzer = TrafficAnalyzer()
        self.violation_detector = ViolationDetector()
        
        # 分析模式设置
        self.enable_traffic_analysis = get_config('behavior_analysis.enable_traffic_analysis', True)
        self.enable_zone_violation = get_config('behavior_analysis.enable_zone_violation', True)
        self.enable_direction_detection = get_config('behavior_analysis.enable_direction_detection', False)
        self.enable_parking_detection = get_config('behavior_analysis.enable_parking_detection', False)
        
        # 历史记录参数
        self.max_history_length = get_config('behavior_analysis.max_history_length', 1000)
        
        # 分析历史
        self.analysis_history = []
        self.last_analysis_time = 0
        
        print("交通监控行为分析管理器已初始化")
    
    def analyze_frame(self, frame: np.ndarray, timestamp: float, 
                     detections: List[Dict], tracking_results: List[Dict] = None) -> Dict:
        """分析视频帧中的交通场景
        
        Args:
            frame: 输入视频帧
            timestamp: 当前时间戳
            detections: YOLO检测结果
            tracking_results: 跟踪结果（可选）
            
        Returns:
            分析结果字典
        """
        # 初始化结果
        results = {
            'timestamp': timestamp,
            'detections': detections,
            'traffic_analysis': None,
            'violations': [],
            'statistics': {}
        }
        
        # 执行交通分析
        if self.enable_traffic_analysis:
            traffic_results = self.traffic_analyzer.analyze_traffic(frame, detections, tracking_results)
            results['traffic_analysis'] = traffic_results
        else:
            traffic_results = None
        
        # 执行违规检测
        violations = []
        if tracking_results:
            violations = self.violation_detector.detect_violations(
                timestamp=timestamp,
                detections=detections,
                tracking_results=tracking_results,
                traffic_results=traffic_results
            )
        results['violations'] = violations
        
        # 收集统计信息
        statistics = {
            'traffic_stats': self.traffic_analyzer.get_statistics(),
            'violation_stats': self.violation_detector.get_violation_statistics()
        }
        results['statistics'] = statistics
        
        # 记录历史
        self._record_history(results)
        
        # 更新最后分析时间
        self.last_analysis_time = timestamp
        
        return results
    
    def _record_history(self, analysis_result: Dict) -> None:
        """记录分析历史
        
        Args:
            analysis_result: 分析结果
        """
        self.analysis_history.append(analysis_result)
        
        # 限制历史记录长度
        if len(self.analysis_history) > self.max_history_length:
            self.analysis_history = self.analysis_history[-self.max_history_length:]
    
    def get_aggregated_statistics(self, time_window: float = 60.0) -> Dict:
        """获取聚合统计信息
        
        Args:
            time_window: 时间窗口（秒）
            
        Returns:
            聚合统计信息
        """
        current_time = time.time()
        window_start = current_time - time_window
        
        # 筛选时间窗口内的历史记录
        window_history = [
            h for h in self.analysis_history 
            if h['timestamp'] >= window_start
        ]
        
        # 初始化聚合统计
        aggregated = {
            'vehicle_counts': {
                'car': 0,
                'bicycle': 0,
                'motorbike': 0,
                'truck': 0,
                'bus': 0,
                'person': 0,
                'total': 0
            },
            'violations_by_type': {
                'zone_violation': 0,
                'wrong_direction': 0,
                'illegal_parking': 0,
                'pedestrian_violation': 0,
                'vehicle_overlap': 0
            },
            'total_violations': 0,
            'time_window': time_window,
            'record_count': len(window_history)
        }
        
        # 聚合统计
        for record in window_history:
            # 聚合车辆统计
            if 'traffic_counts' in record['statistics']:
                traffic_counts = record['statistics']['traffic_counts']
                for vehicle_type, count in traffic_counts.items():
                    if vehicle_type in aggregated['vehicle_counts']:
                        aggregated['vehicle_counts'][vehicle_type] += count
                        aggregated['vehicle_counts']['total'] += count
            
            # 聚合违规统计
            for violation in record['violations']:
                violation_type = violation['type']
                if violation_type in aggregated['violations_by_type']:
                    aggregated['violations_by_type'][violation_type] += 1
                aggregated['total_violations'] += 1
        
        return aggregated
    
    def set_analysis_modes(self, enable_traffic: bool = None, 
                          enable_zone: bool = None,
                          enable_direction: bool = None,
                          enable_parking: bool = None,
                          enable_pedestrian: bool = None,
                          enable_overlap: bool = None) -> None:
        """设置分析模式
        
        Args:
            enable_traffic: 是否启用交通分析
            enable_zone: 是否启用区域违规检测
            enable_direction: 是否启用方向违规检测
            enable_parking: 是否启用停车违规检测
            enable_pedestrian: 是否启用行人违规检测
            enable_overlap: 是否启用车辆重叠检测
        """
        if enable_traffic is not None:
            self.enable_traffic_analysis = enable_traffic
        if enable_zone is not None:
            self.enable_zone_violation = enable_zone
        if enable_direction is not None:
            self.enable_direction_detection = enable_direction
        if enable_parking is not None:
            self.enable_parking_detection = enable_parking
        if enable_pedestrian is not None:
            self.enable_pedestrian_violation = enable_pedestrian
        if enable_overlap is not None:
            self.enable_vehicle_overlap = enable_overlap
    
    def set_traffic_analysis_params(self, **kwargs) -> None:
        """设置交通分析参数
        
        Args:
            **kwargs: 交通分析参数
        """
        # 阈值参数
        if 'confidence_threshold' in kwargs:
            self.traffic_analyzer.set_confidence_threshold(kwargs.get('confidence_threshold'))
        
        # 交通类别权重参数
        if 'vehicle_weights' in kwargs:
            self.traffic_analyzer.set_vehicle_weights(kwargs.get('vehicle_weights'))
        
        # 跟踪参数
        if 'tracking_params' in kwargs:
            self.traffic_analyzer.set_tracking_params(kwargs.get('tracking_params'))
        
        # 交通分析参数已在初始化时设置
    
    def set_violation_detection_params(self, **kwargs) -> None:
        """设置违规检测参数
        
        Args:
            **kwargs: 违规检测参数
        """
        # 违规基本参数
        if 'min_duration' in kwargs or 'max_interval' in kwargs or 'zone_cooldown' in kwargs:
            self.violation_detector.set_violation_params(
                min_duration=kwargs.get('min_duration'),
                max_interval=kwargs.get('max_interval'),
                zone_cooldown=kwargs.get('zone_cooldown')
            )
        
        # 违规权重
        if 'weights' in kwargs:
            self.violation_detector.set_violation_weights(kwargs['weights'])
    
    def add_restricted_zone(self, zone_id: str, name: str, points: List[Tuple],
                           prohibited_classes: List[str] = None) -> None:
        """添加禁行区域
        
        Args:
            zone_id: 区域ID
            name: 区域名称
            points: 区域多边形点列表
            prohibited_classes: 禁止进入的类别列表
        """
        self.violation_detector.add_restricted_zone(zone_id, name, points, prohibited_classes)
    
    def add_bicycle_lane(self, lane_id: str, name: str, points: List[Tuple]) -> None:
        """添加自行车道
        
        Args:
            lane_id: 车道ID
            name: 车道名称
            points: 车道多边形点列表
        """
        self.violation_detector.add_bicycle_lane(lane_id, name, points)
    
    def add_vehicle_direction(self, lane_id: str, lane_points: List[Tuple],
                             allowed_direction: Tuple, allowed_classes: List[str]) -> None:
        """添加车辆行驶方向限制
        
        Args:
            lane_id: 车道ID
            lane_points: 车道多边形点列表
            allowed_direction: 允许的行驶方向向量
            allowed_classes: 允许的车辆类别
        """
        self.violation_detector.add_vehicle_direction(
            lane_id, lane_points, allowed_direction, allowed_classes
        )
    
    def clear_all_zones(self) -> None:
        """清除所有区域设置"""
        self.violation_detector.clear_all_zones()
    
    def get_analysis_history(self, limit: int = None) -> List[Dict]:
        """获取分析历史
        
        Args:
            limit: 返回记录数量限制
            
        Returns:
            分析历史记录列表
        """
        if limit is None:
            return self.analysis_history.copy()
        return self.analysis_history[-limit:].copy()
    
    def clear_analysis_history(self) -> None:
        """清除分析历史"""
        self.analysis_history.clear()
    
    def get_active_violations(self) -> Dict:
        """获取当前活跃的违规
        
        Returns:
            活跃违规字典
        """
        return self.violation_detector.get_active_violations()
    
    def visualize_results(self, frame: np.ndarray, analysis_results: Dict) -> np.ndarray:
        """可视化分析结果
        
        Args:
            frame: 输入视频帧
            analysis_results: 分析结果
            
        Returns:
            可视化后的帧
        """
        visualized_frame = frame.copy()
        
        # 可视化交通分析结果
        if analysis_results['traffic_analysis']:
            visualized_frame = self._visualize_traffic_analysis(
                visualized_frame,
                analysis_results['traffic_analysis']
            )
        
        # 可视化违规结果
        if analysis_results['violations']:
            visualized_frame = self._visualize_violations(
                visualized_frame,
                analysis_results['violations']
            )
        
        # 可视化统计信息
        visualized_frame = self._visualize_statistics(
            visualized_frame,
            analysis_results['statistics']
        )
        
        return visualized_frame
    
    def _visualize_traffic_analysis(self, frame: np.ndarray, traffic_results: Dict) -> np.ndarray:
        """可视化交通分析结果
        
        Args:
            frame: 输入视频帧
            traffic_results: 交通分析结果
            
        Returns:
            可视化后的视频帧
        """
        # 复制帧以避免修改原始数据
        vis_frame = frame.copy()
        
        # 设置各类别颜色
        class_colors = {
            'car': (0, 0, 255),       # 红色
            'bicycle': (0, 255, 255), # 黄色
            'motorbike': (255, 0, 0), # 蓝色
            'person': (0, 255, 0),    # 绿色
            'truck': (255, 255, 0),   # 青色
            'bus': (255, 0, 255)      # 洋红色
        }
        
        # 绘制交通检测结果
        if 'detected_objects' in traffic_results:
            for detection in traffic_results['detected_objects']:
                # 获取边界框坐标
                bbox = detection.get('bbox', [])
                if len(bbox) >= 4:
                    x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
                    
                    # 获取类别
                    obj_class = detection.get('class', 'unknown')
                    color = class_colors.get(obj_class, (128, 128, 128))  # 默认灰色
                    
                    # 绘制边界框
                    cv2.rectangle(vis_frame, (x1, y1), (x2, y2), color, 2)
                    
                    # 添加标签
                    label = f"{obj_class} {detection.get('confidence', 0):.2f}"
                    cv2.putText(vis_frame, label, (x1, y1 - 10), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
        
        # 绘制车辆计数信息
        if 'traffic_counts' in traffic_results:
            count_text = "车辆统计: "
            for vehicle_type, count in traffic_results['traffic_counts'].items():
                if count > 0:
                    count_text += f"{vehicle_type}: {count} "
            
            # 在左上角显示统计信息
            cv2.putText(vis_frame, count_text, (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        return vis_frame
    
    def _visualize_violations(self, frame: np.ndarray, 
                             violations: List[Dict]) -> np.ndarray:
        """可视化违规结果
        
        Args:
            frame: 输入帧
            violations: 违规结果
            
        Returns:
            可视化后的帧
        """
        # 违规类型颜色映射
        violation_colors = {
            'zone_violation': (255, 0, 255),  # 洋红色
            'wrong_direction': (255, 165, 0), # 橙色
            'illegal_parking': (255, 255, 0)  # 黄色
        }
        
        # 违规类型标签映射
        violation_labels = {
            'zone_violation': "Zone Violation",
            'wrong_direction': "Wrong Direction",
            'illegal_parking': "Illegal Parking"
        }
        
        for violation in violations:
            violation_type = violation['type']
            color = violation_colors.get(violation_type, (128, 128, 128))
            label = violation_labels.get(violation_type, violation_type)
            
            # 绘制违规区域
            location = violation['location']
            x1, y1, x2, y2 = location['x1'], location['y1'], location['x2'], location['y2']
            
            # 绘制闪烁效果的边界框
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            cv2.rectangle(frame, (x1-1, y1-1), (x2+1, y2+1), color, 1)
            
            # 绘制违规标签
            violation_text = f"VIOLATION: {label}"
            cv2.putText(
                frame,
                violation_text,
                (x1, max(y1 - 30, 10)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                2
            )
            
            # 绘制违规严重程度
            severity_text = f"Severity: {violation['severity']:.1f}"
            cv2.putText(
                frame,
                severity_text,
                (x1, max(y1 - 15, 25)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4,
                color,
                1
            )
        
        return frame
    
    def _visualize_statistics(self, frame: np.ndarray, statistics: Dict) -> np.ndarray:
        """可视化统计信息
        
        Args:
            frame: 输入帧
            statistics: 统计信息
            
        Returns:
            可视化后的帧
        """
        # 在左上角绘制统计信息
        stats_text = [
            f"Total Violations: {statistics['violation_stats']['total']}"
        ]
        
        # 交通统计信息
        if 'traffic_stats' in statistics and 'vehicle_counts' in statistics['traffic_stats']:
            vehicle_counts = statistics['traffic_stats']['vehicle_counts']
            stats_text.extend([
                f"Cars: {vehicle_counts.get('car', 0)}",
                f"Bicycles: {vehicle_counts.get('bicycle', 0)}",
                f"Motorcycles: {vehicle_counts.get('motorbike', 0)}",
                f"Trucks: {vehicle_counts.get('truck', 0)}",
                f"Pedestrians: {vehicle_counts.get('person', 0)}",
                f"Total Vehicles: {vehicle_counts.get('total', 0)}"
            ])
        
        y_offset = frame.shape[0] - 10
        for text in reversed(stats_text):
            y_offset -= 20
            cv2.putText(
                frame,
                text,
                (10, y_offset),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4,
                (0, 255, 0),
                1
            )
        
        return frame

# 创建全局行为分析管理器实例
analysis_manager = AnalysisManager()

# 导出常用方法
def get_analysis_manager() -> AnalysisManager:
    """获取行为分析管理器实例的快捷方法"""
    return analysis_manager